# Function to make sure 
f <- function(dat)
{
  ifelse(is.null(dat), NA, dat)
}

# Function to get cap rates -----------------------------------------------
Generate_Contract_Specs <- function(logR,
                                    r,
                                    mu_d_RN,
                                    MRP,
                                    Cost_grid = c(0.002, 0.003, 0.005, 0.007))
{
  
  # Testing
  # input <- list()
  # input$age = 35
  # input$retirement_age = 65
  # input$income = 1000
  # year_retire = input$retirement_age - input$age + 2022
  # logR = get_log_returns()
  # r = logR$r # constant annual interest rate, compounded continuously.
  # Cost_grid = c(0.002, 0.003, 0.005, 0.007) # cost ('fee_TDF' or 'c_bar_RILA').
  # MRP = logR$MRP
  # mu_d_RN = logR$mu_d_RN
  
  # Other vars
  mu_d = mu_d_RN + MRP # mean of diffusion process (under real-world measure).
  C_num = length(Cost_grid)  
  
  
  
  
  # Contract specs ----------------------------------------------------------
  
  # number of contracts used in experiment.
  num_con = 9                    
  
  con_type_vec <- c()
  con_type_vec[1:3] = "TDF"
  con_type_vec[4:6] = "Buffer"
  con_type_vec[7:9] = "Floor"
  
  # contains 'alpha' for TDF, 'B' for buffer-RILA and 'F' for floor-RILA.
  con_spec_vec = matrix(0, 50, num_con)
  
  # TDFs:
  con_spec_vec[1:20, 1] = 1       
  con_spec_vec[21:50, 1] = 0.95 - 0.01 * c(1:30)
  con_spec_vec[, 2] = con_spec_vec[, 1] - 0.10
  con_spec_vec[, 3] = con_spec_vec[, 1] - 0.20
  
  # Buffer RILAs:
  con_spec_vec[1:20, 4] = 0     
  con_spec_vec[21:50, 4] = 1/300 * c(1:30)
  con_spec_vec[, 5] = con_spec_vec[, 4] + 0.025
  con_spec_vec[, 6] = con_spec_vec[, 4] + 0.050
  
  # Floor RILAs:
  con_spec_vec[1:20, 7] = 0.25  	
  con_spec_vec[21:50, 7] = 0.25 - 1.9/300 * c(1:30)
  con_spec_vec[, 8] = con_spec_vec[, 7] - 0.035
  con_spec_vec[, 9] = con_spec_vec[, 7] - 0.055
  
  
  # Find Cap Rates of RILA contracts: ---------------------------------------
  
  # contains cap rates for RILA contracts.
  con_cap_vec = array(0, dim = c(50, C_num, num_con))     
  
  # Go ahead and exponentiate log returns
  R_S = exp(logR) - 1 
  # R_S = exp(logR[[1]][,1]) - 1 
  
  # for(t in 1:50){
  #   
  #   for(c in 1:C_num){
  #     
  #     for(j in 1:num_con){
  #       
  #       if(con_type_vec[j] == "TDF"){
  #         # do nothing
  #         
  #       } else if(con_type_vec[j] == "Buffer") {
  #         con_cap_vec[t,c,j] = optimize(ExpLoss, 
  #                                      interval = c(0, 1),
  #                                      D_level = con_spec_vec[t,j],
  #                                      D_type = "Buffer",
  #                                      r = r,
  #                                      tau = 1,
  #                                      fee = 0, 
  #                                      c_bar = Cost_grid[c],
  #                                      R_S = R_S)$minimum 
  #         
  #       } else if(con_type_vec[j] == "Floor") {
  #         con_cap_vec[t,c,j] = optimize(ExpLoss, 
  #                                       interval = c(0, 1),
  #                                       D_level = con_spec_vec[t,j],
  #                                       D_type = "Buffer",
  #                                       r = r,
  #                                       tau = 1,
  #                                       fee = 0, 
  #                                       c_bar = Cost_grid[c],
  #                                       R_S = R_S)$minimum 
  #         
  #       } else {
  #         print('Contract type unspecified.')
  #         
  #       }
  #       
  #     }
  #     
  #   }
  #   print(t*c*j/(50 * C_num * num_con))
  # }
  
  load(file = "./Code/Data/ContractSpecs.RData")
  
  output <- list(
    Cost_grid = ContractSpecs$Cost.grid,
    con_type_vec = con_type_vec,
    con_spec_vec = ContractSpecs$con.spec.vec,
    con_cap_vec = ContractSpecs$con.cap.vec)
  
  return(output)
}


# Function to help find fair RILA cap rates -------------------------------
ExpLoss <- function(Cap,
                    D_level,
                    D_type,
                    r,
                    tau,
                    fee,
                    c_bar,
                    R_S)
{
  # Cap = 1
  # D_level = con_spec_vec[t,j]
  # D_type = "Buffer"
  # r = r
  # tau = 1
  # fee = 0
  # c_bar = c_bar_RILA
  # logR = logR
  # A_t = 1000
  
  # simulated annual net returns of market index.
  A_t = 1000
  
  if(D_type == "Buffer"){
    
    # vector (N_sim-by-1) of credited return to RILA account.
    R_A = pmin(R_S + D_level, 0) * (R_S < 0) + pmin(R_S, Cap) * (R_S > 0) 
    
  } else if(D_type == "Floor"){
    
    # vector (N_sim-by-1) of credited return to RILA account.
    R_A = pmax(R_S , -D_level) * (R_S < 0) + pmin(R_S, Cap) * (R_S > 0) 
    
  } else {
    print('Error. Downside protection type unspecified.')
    asdfadsf    # to produce error
  }
  
  A_tplustau = A_t * exp(-fee * tau) * (1 + R_A)
  L = A_tplustau * exp(-r * tau) + A_t * c_bar * tau - A_t
  EL = mean(L)
  
  return(EL^2)
}


# Function to get log returns ---------------------------------------------

# Output: logR(i,t) = simulated log return of market index under path i in year t. 

get_log_returns <- function(regime = "MJD",
                            infl = 0.02,
                            MRP = 0.06,
                            delta = 0,
                            max_num_cont = 50,
                            N_sim = 1e+5,
                            r = 0.03)
{
  max_num_cont = 50  # maximum number of future QRP contributions
  regime = "MJD" # Financial market framework
  r = 0.03           # constant annual interest rate, compounded continuously.
  infl = 0.02        # annual rate of inflation.
  MRP = 0.06         # market risk premium (for investment in S&P 500, = 100% volatility).
  delta = 0          # dividend yield (assume =0 to make RILAs better comparable to TDFs).
  N_sim = 1e+5      #  number of simulated paths.
  
  if(regime == "MJD"){
    lambda = 0.20      # likelihood of jump.
    mu_d_RN = 0.0735   # mean of diffusion process (under risk-neutral measure).
    sigma_d = 0.14     # volatility of diffusion process.
    mu_j = -0.25       # mean of each jump.
    sigma_j = 0.10     # volatility of magnitude of each jump.
  } else if(regime == "BS"){
    lambda = 0         # likelihood of jump.
    mu_d_RN = 0.03     # mean of diffusion process (under risk-neutral measure).
    sigma_d = 0.18     # volatility of diffusion process.
    mu_j = 0           # mean of each jump.
    sigma_j = 0.01     # volatility of magnitude of each jump.
  } else {
    print('Regime not specified correctly.')
    asdadsf
  }
  
  mu_d = mu_d_RN + MRP   # mean of diffusion process (under real-world measure).
  
  # simulate log returns of index for 1 year:
  num_jumps <- matrix(0, N_sim, max_num_cont)
  for(t in 1:max_num_cont){
    num_jumps[,t] = rpois(N_sim, lambda)
  }
  
  # vector of annual log returns, following MJD model.
  logR <- matrix(0, N_sim, max_num_cont)
 
  for(t in 1:max_num_cont){
  logR[,t] = (mu_d - sigma_d^2/2)*1 + 
    num_jumps[,t] * mu_j + 
    sqrt( sigma_d^2/2 * 1 + num_jumps[,t] * sigma_j^2 ) * rnorm(n = N_sim, mean = 0, sd = 1)   
}
  load(file = "./Code/Data/LogReturns.RData")
  output <- list(
    logR = LogReturns$logR,
    # logR = logR,
    r = r,
    infl = infl,
    N_sim = N_sim,
    mu_d_RN = mu_d_RN,
    MRP = MRP)
  
  return(output)
}


# Get simulated returns ---------------------------------------------------
Get_simul_stats_fun1 <- function(con_type,
                                 con_spec,
                                 con_cap,
                                 fee_TDF,
                                 fee_RILA,
                                 cont_t,
                                 logR,
                                 A_0,
                                 r,
                                 ytd,
                                 infl)
{
  # alpha_vec = vector of annual equity allocation percentages in TDF.
  # fee = annual fee rate of TDF.
  # A_0 = investor's starting QRP account value. 
  # I_1 = investor's first contribution (made at end of first year).
  # g_I = annual growth rate of contributions.
  # logR = matrix of annual log returns of index (--> dimensions: paths & time).  
  # cont_t = vector of annual contribution amounts (--> cont_t(t) = amount contributed at time t).
  
  
  N_sim = length(logR[,1])       # number of simulated paths.
  num_cont = length(logR[1,])    # number of future contributions made by investor.
  
  
  
  # Update TDF account value A_t --------------------------------------------
  
  A_t = matrix(0, N_sim, num_cont+1)
  A_t[,1] = A_0
  
  for(t in 1:num_cont){
    
    R_S = exp(logR[,t])-1    # annual net return of stock index.
    
    if(con_type == "TDF"){
      R_A = ( con_spec[t] * exp(logR[,t]) + (1-con_spec[t]) * exp(r) ) * exp(-fee_TDF)
    } else if(con_type == "Buffer") {
      B = con_spec[t]
      Cap = con_cap[t]
      R_A = pmin(R_S+B, 0) * (R_S < 0) + pmin(R_S, Cap) * (R_S > 0)      # vector (N_sim-by-1) of credited return to RILA account.
    } else if(con_type == "Floor"){
      F = con_spec[t]
      Cap = con_cap[t]
      R_A = pmax(R_S, -F) * (R_S < 0) + pmin(R_S, Cap) * (R_S > 0)     # vector (N_sim-by-1) of credited return to RILA account.
    } else {
      disp('Contract type undefined.')
    }
    
    I_t = cont_t[t]  # PH's time-t contribution to QRP (no contribution at retirement time T)
    if(con_type == "TDF") {
      A_t[,t+1] = A_t[,t] * R_A + I_t
    } else {
      # RILA contract
      A_t[,t+1] = A_t[,t] * (1 + R_A) * exp(-fee_RILA) + I_t
    }
  }
  
  A_T = A_t[,num_cont+1]   # vector (size N_sim-by-1) of simulated QRP account values at retirement.
  
  # Find statistics of terminal account value
  
  # convert terminal account value to whole life annuity, inflation-adjusted (--> priced at constant rate 'r' and assuming 'ytd' years to death; payments made at beg. of each month). 
  r_net = r - infl
  i_m = 12*(exp(r_net/12)-1)             # i^{(m)}, the nominal interest rate compounded monthly.
  annuity_factor = (1-exp(-r_net*ytd))/i_m
  ret_inc = A_T/12/annuity_factor * (1+infl)^(-num_cont)       # monthly retirement income (in today's $$).
  
  # EV = mean(ret_inc);
  V_bins = matrix(0,1,12) # CTE(1) = min.  CTE(12) = max. CTE(2:11) = average of lowest 10% of outcomes, second-lowest 10%, ... , highest 10%.
  
  V_bins[1] = mean(ret_inc[which(ret_inc <= quantile(ret_inc, .1))])
  V_bins[12] = mean(ret_inc[which(ret_inc >= quantile(ret_inc, .99))])
  
  for(i in 1:10){
    V_bins[i+1] = mean(ret_inc[which((ret_inc >= quantile(ret_inc,.10*(i-1))) &
                                        (ret_inc <= quantile(ret_inc,.10*i))  )] )
  }
  
  return(V_bins)
}


Get_V_bins <- function(CurrentAge,
                       CurrentAccountValue,
                       InitialContributionAmount,
                       RetirementAge,
                       AnnualCost,
                       logR,
                       r,
                       infl,
                       Cost_grid,
                       con_type_vec,
                       con_spec_vec,
                       con_cap_vec)
{
  
  # CurrentAge = 35
  # CurrentAccountValue = 1000
  # InitialContributionAmount = 5000
  # RetirementAge = 65
  # AnnualCost = .003
  # logR = lr$logR
  # r = lr$r
  # infl = lr$infl
  # Cost_grid = cs$Cost_grid
  # con_type_vec = cs$con_type_vec
  # con_spec_vec = cs$con_spec_vec
  # con_cap_vec = cs$con_cap_vec
  
  cost_loc = which(Cost_grid == AnnualCost)   # location of alternative cost on Cost_grid.
  
  if(length(cost_loc) != 1) print('Error. Alternative cost not on grid.')
  
  A_0 = CurrentAccountValue
  cont_1 = InitialContributionAmount
  YTR = RetirementAge - CurrentAge      # Years to retirement (= # of future QRP contributions).
  ytd = 85 - RetirementAge              # Years from retirement to death (--> used only to price life annuity).
  max_num_cont = length(con_spec_vec[,1])  # maximum # of future QRP contributions.
  fee_TDF = AnnualCost
  
  g_I = 0.03     # annual rate of salary (and QRP contribution) increase.
  fee_RILA = 0
  num_con = 9
  
  ## Determine simulation statistics
  
  V_bins = matrix(0, num_con, 12)
  
  for(j in 1:num_con){               # cycle through all 9 contracts.
    cont_t = cont_1*(1+g_I)^(0:(YTR-1))               # vector of QRP contribution amounts each year (--> cont_t(t) = amount contributed at time t).
    con_type = con_type_vec[j]
    con_spec = con_spec_vec[(max_num_cont-YTR+1):max_num_cont, j]         # NOTE: if individual has 35 years left till retirement, we should use the contract specs from years 16 through 50.
    con_cap = con_cap_vec[(max_num_cont-YTR+1):max_num_cont,cost_loc, j]
    
    V_bins[j,] = Get_simul_stats_fun1(con_type = con_type,
                                      con_spec = con_spec,
                                      con_cap = con_cap,
                                      fee_TDF = fee_TDF,
                                      fee_RILA = fee_RILA,
                                      cont_t = cont_t,
                                      logR = logR[,1:YTR],
                                      A_0 = A_0,
                                      r = r,
                                      ytd = ytd,
                                      infl = infl)
  }
  return(V_bins)
}
